package com.swiftops.common.util;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

public class TreeUtils {

    public static <T> List<T> buildTree(List<T> flatList, Function<T, String> idGetter,
                                        Function<T, String> parentIdGetter, String rootId) {

        Map<String, T> idToNodeMap = new HashMap<>();
        List<T> roots = new ArrayList<>();

        // first pass: create nodes and map them by ID
        for (T node : flatList) {
            String id = idGetter.apply(node);
            idToNodeMap.put(id, node);
        }

        // second pass: link nodes and create root list
        for (T node : flatList) {
            String parentId = parentIdGetter.apply(node);
            if (parentId == null || parentId.equals(rootId)) {
                roots.add(node);
            } else {
                T parent = idToNodeMap.get(parentId);
                if (parent != null) {
                    List<T> children = (List<T>) getOrCreateChildren(parent);
                    children.add(node);
                }
            }
        }

        return roots;
    }

    private static Object getOrCreateChildren(Object node) {
        try {
            Method method = node.getClass().getMethod("getChildren");
            List<Object> children = (List<Object>) method.invoke(node);
            if (children == null) {
                children = new ArrayList<>();
                method = node.getClass().getMethod("setChildren", List.class);
                method.invoke(node, children);
            }
            return children;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
